from __future__ import division


class CubicSpline(object):
    def __init__(self, grid):
        self.grid = grid
        self.calculate_coeffs()

    def calculate_coeffs(self):
        n = len(self.grid) - 1
        A = [
                [self.hi(i)/3, 2/3 * (self.hi(i) + self.hi(i+1)), self.hi(i+1)/3,
                    (self.grid.yvalues[i+1] - self.grid.yvalues[i]) / self.hi(i+1) -
                    (self.grid.yvalues[i] - self.grid.yvalues[i-1]) / self.hi(i)]
                for i in xrange(1, n)
                ]
        A[0][0] = 0
        A[-1][-2] = 0
        self.a_coeffs = [0]*(n+1)
        self.b_coeffs = [0]*(n+1)
        self.c_coeffs = [0] + self.solve_tridiagonal_system(A)
        self.c_coeffs.append(0)
        self.d_coeffs = [0]*(n+1)
        for i in xrange(1, n+1):
            self.a_coeffs[i] = self.grid.yvalues[i-1]
            self.b_coeffs[i] = (self.grid.yvalues[i] - self.grid.yvalues[i-1]) / self.hi(i) -\
                    self.c_coeffs[i] * self.hi(i) -\
                    (self.c_coeffs[i+1] - self.c_coeffs[i]) * self.hi(i) / 3
            self.d_coeffs[i] = (self.c_coeffs[i+1] - self.c_coeffs[i]) / (3 * self.hi(i))

    def hi(self, i):
        return self.grid.xvalues[i] - self.grid.xvalues[i-1]

    def interpolate(self, x):
        if x < self.grid.xvalues[0] or x > self.grid.xvalues[-1]:
            raise RuntimeError('x is out of the interval')
        for i in range(len(self.grid)):
            if x < self.grid.xvalues[i]:
                return self.a_coeffs[i] +\
                        self.b_coeffs[i] * (x - self.grid.xvalues[i-1]) +\
                        self.c_coeffs[i] * (x - self.grid.xvalues[i-1])**2 +\
                        self.d_coeffs[i] * (x - self.grid.xvalues[i-1])**3

    def solve_tridiagonal_system(self, A):
        n = len(A)
        a_coeffs = [0]*(n+1)
        b_coeffs = [0]*(n+1)
        for i in xrange(1, n+1):
            a_coeffs[i] = - A[i-1][2] / (A[i-1][0]*a_coeffs[i-1] + A[i-1][1])
            b_coeffs[i] = (A[i-1][3] - A[i-1][0]*b_coeffs[i-1]) / (A[i-1][0]*a_coeffs[i-1] + A[i-1][1])
        x_values = [0]*(n+1)
        x_values[n] = b_coeffs[n]
        for i in xrange(n-1, 0, -1):
            x_values[i] = a_coeffs[i] * x_values[i+1] + b_coeffs[i]
        return x_values
